import asyncio
from agents import Agent, Runner, OpenAIProvider, RunResult, RunConfig, function_tool, TResponseInputItem, set_tracing_export_api_key, set_tracing_disabled
from openai import AsyncOpenAI
from dataclasses import dataclass
from typing import Literal, List, Dict
import os

# 实验指定的API配置
BASE_URL = "https://api.chatfire.cn/v1"
API_KEY = "sk-UcrTOYgPf03qqxdGspmJtaUmP7OTkkzuxsRHWPpMSxKXK8IN"
os.environ["OPENAI_API_KEY"] = API_KEY

# 配置Tracing客户端使用相同的API密钥
set_tracing_export_api_key(API_KEY)
set_tracing_disabled(True)

# ------------------------- 示例1：单Agent函数调用 -------------------------
@function_tool
def get_weather(city: str) -> str:
    """天气查询工具函数"""
    return f"The weather in {city} is sunny."

provider = OpenAIProvider(
        openai_client=AsyncOpenAI(
            base_url=BASE_URL,
            api_key=API_KEY,
            timeout=30.0  # 增加超时时间
        ),
        use_responses=False
    )

async def demo_agent_function_call():
    """演示单个Agent调用工具函数的流程"""
    agent = Agent(
        name="WeatherAgent",
        instructions="你是一个能通过工具查询天气的智能助手",
        tools=[get_weather],
        model="gpt-4o"
    )
    result = await Runner.run(
        agent,
        input="What's the weather in Beijing?",
        run_config=RunConfig(model_provider=provider)
    )
    print("[Agent函数调用结果]")
    print(result.final_output)


# ------------------------- 示例2：MultiAgent工作流 -------------------------
@dataclass
class EvaluationFeedback:
    """评估反馈数据结构"""
    score: Literal["pass", "needs_improvement", "fail"]
    feedback: str


async def demo_multi_agent_workflow():
    """演示多个Agent协同工作的流程"""
    # 定义故事大纲生成Agent
    story_generator = Agent(
        name="StoryOutlineGenerator",
        model="gpt-4o",
        instructions=(
            "根据用户输入生成简短故事大纲\n"
            "如果有反馈，请基于反馈优化大纲"
        )
    )
    # 定义评估Agent
    evaluator = Agent(
        name="StoryEvaluator",
        model="gpt-4o",
        instructions=(
            "评估故事大纲是否足够好\n"
            "若不够好，提供具体改进建议\n"
            "首次评估不得直接通过"
        ),
        output_type=EvaluationFeedback
    )

    # 初始化用户输入
    story_type = input("请输入想生成的故事类型（如：科幻冒险故事）: ")
    input_items: List[TResponseInputItem] = [{"content": story_type, "role": "user"}]
    latest_outline = None

    while True:
        # 1. 生成故事大纲
        gen_result = await Runner.run(
            story_generator,
            input_items,
            run_config=RunConfig(model_provider=provider)
        )
        input_items = gen_result.to_input_list()
        latest_outline = gen_result.final_output_as(str)
        print("\n[生成的故事大纲]")
        print(latest_outline)

        # 2. 评估大纲
        eval_result = await Runner.run(
            evaluator,
            input_items,
            run_config=RunConfig(model_provider=provider)
        )
        feedback: EvaluationFeedback = eval_result.final_output
        print(f"[评估结果] 分数: {feedback.score}")

        # 3. 根据反馈迭代
        if feedback.score == "pass":
            print("大纲已通过评估，结束流程")
            break
        print(f"[改进建议] {feedback.feedback}")
        input_items.append({
            "content": f"Feedback: {feedback.feedback}",
            "role": "user"
        })
        print("基于反馈重新生成大纲...")
        break

    print(f"\n[最终确定的故事大纲]\n{latest_outline}")


# ------------------------- 主执行流程 -------------------------
async def main():
    """主函数：按顺序执行两个演示案例"""
    print("===== OpenAI Agent实验演示 =====")

    # 执行单Agent演示
    print("\n--- 演示1：单Agent函数调用 ---")
    await demo_agent_function_call()

    # 等待用户确认后执行MultiAgent演示
    input("\n按Enter键继续执行MultiAgent工作流演示...")

    # 执行MultiAgent演示
    print("\n--- 演示2：MultiAgent工作流协同 ---")
    await demo_multi_agent_workflow()


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except Exception as e:
        print(f"执行过程中发生错误: {str(e)}")